package com.zzyl.common.ai;


import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.concurrent.TimeUnit;

@Component
public class AIModelInvoker {

    @Autowired
    private BaiduAIProperties baiduAIProperties;


    public static final OkHttpClient HTTP_CLIENT = new OkHttpClient().newBuilder().readTimeout(300, TimeUnit.SECONDS).build();

    /****
     * 远程调用
     */
    public String qianfanInvoker(String prompt) throws IOException {
        MediaType mediaType = MediaType.parse("application/json");
        ObjectMapper mapper = new ObjectMapper();
        ObjectNode root = mapper.createObjectNode();

        root.put("model", baiduAIProperties.getQianfanModel());
        root.set("messages", mapper.readTree(prompt));
        root.put("stream", false);
        String json = mapper.writeValueAsString(root);
        RequestBody body = RequestBody.create(mediaType, json);
        Request request = new Request.Builder()
                .url("https://qianfan.baidubce.com/v2/chat/completions")
                .method("POST", body)
                .addHeader("Content-Type", "application/json")
                .addHeader("appid", baiduAIProperties.getAppid())
                .addHeader("Authorization", baiduAIProperties.getAccessKey())
                .build();
        Response response = HTTP_CLIENT.newCall(request).execute();
        return response.body().string();
    }
}
